AIR CLI Integration: Port code_source Snapshotting#5897
Conversation
First slice of the code_source snapshot port (Go equivalent of the Python CLI's cli/utils/snapshot.py). Adds the snapshot subpackage with the pure cache-key function and a golden-fixture harness: - cachekey.go: ComputeCacheKey + PackagingVersion="v1", ported verbatim from compute_snapshot_cache_key so both CLIs agree on the key algorithm. - testdata/cache_keys.json: golden keys captured from the exact Python algorithm over the local-only matrix (commit + include_paths permutations). - cachekey_test.go: byte-for-byte golden table test plus a properties test pinning the normalization (trim + sort, no dedup, nil == empty). Co-authored-by: Isaac
…only) Continues the code_source snapshot port. Flattens the snapshot files into the aircmd package (matching the rest of experimental/air/cmd; an organizational layer can land later) and adds Phase 2/3: - snapshot_git.go: local, no-network git introspection ported from cli/utils/git_state.py — is-repo, HEAD/branch SHA, dirty check (whole-tree and path-scoped), commit-exists, include-path validation. Concrete gitRepo shelling out with arg slices (no shell string); tested against real temp repos, matching the Python tests. - snapshot_resolve.go: the mode/ref decision tree (cli_entrypoint ~1541-1722), local-only. git.commit (must exist locally) / git.branch (local HEAD) -> git_archive; no ref or non-git dir -> plain_tar. Dirty check runs once and is threaded into the plan; dirty + git.branch is a hard error. - snapshot_package.go: git archive / plain tar builders + .gitignore->exclude parsing ported from cli/utils/snapshot.py. Shells out for parity; preserves the top-level dir name (--prefix / -C parent) the remote extraction depends on, and the .git / macOS AppleDouble exclusions. - runconfig.go: reject a truthy git.remote (remote-fetch path deprecated) with an actionable error pointing to git.commit; remote: false stays valid. The remote-fetch flavor (fetch_branch_sha, partial clone, remote HEAD) is intentionally not ported: the snapshot archives the local copy only. Co-authored-by: Isaac
Completes the code_source snapshot port (Phases 4-5): git provenance sidecars, the top-level orchestrator, upload via libs/filer, and the submit wiring. - snapshot_gitstate.go: git_state.json (schema_version 1, Python field names, null-for-absent encoding) + git_diff.patch capture (dirty-only, 1MB cap, 5s timeout). Best-effort — a failure warns and continues, never fails submit. Merge-base/repo-url read local refs only (no fetch), matching the local-only design. - snapshot.go: runSnapshot orchestrator (resolve -> package+upload -> sidecars). git_archive is cache-keyed and a workspace/Volume existence check skips packaging+upload on a hit; plain_tar is timestamp-named. Uploads through libs/filer (WorkspaceFilesClient, or FilesClient for a remote_volume). Also ports root_path resolution (project_root/, ~, env vars). - snapshot_git.go: add currentBranch / remoteURL / mergeBaseWithUpstream and a runBytes variant for the raw diff bytes. - runsubmit.go: drop the code_source rejection; package the snapshot and attach code_source_path / git_state_path / git_diff_path to the ai_runtime_task. Tested end to end against the in-process fake workspace, including the cache-hit skip and dirty-diff sidecar paths. Co-authored-by: Isaac
Phase 6 of the snapshot port — acceptance tests over the user-visible surface. - experimental/air/run: add dry-run cases for the new code_source config surface — a valid snapshot config passes validation, and a git.remote config is rejected with the actionable error. - experimental/air/run-submit: a full (non-dry-run) submit that packages a git code_source, uploads the tarball + git_state sidecar, and asserts code_source_path / git_state_path land on the ai_runtime_task payload. The commit is date-pinned so the cache-key-derived tarball name is stable, and RecordRequests captures the submitted payload. Co-authored-by: Isaac
Integration test reportCommit: 2964af6
8 interesting tests: 4 RECOVERED, 4 SKIP
Top 10 slowest tests (at least 2 minutes):
|
|
isaac-bot review this PR |
ben-hansen-db
left a comment
There was a problem hiding this comment.
exercised git-archive, plain-tar (dirty, with git_diff.patch), cache-hit reuse, and the dirty+git.branch guard; inspected the uploaded tarball and git_state.json.
great, so commit and branch are tested?
What's the speed like for universe filtered to research folder? Can we also verify that .gitignore is being followed?
| @@ -0,0 +1,2 @@ | |||
| bundle: | |||
There was a problem hiding this comment.
something to consider is a "how to review this pr" section in your readme which can detail what are the important aspects of the pr to consider.
I'd also love to see screenshots from this working e2e
There was a problem hiding this comment.
thanks for this note, i will add in a 'how to review this pr' in future PRs. I have also added screenshots from the e2e tests I ran.
Regarding the speed for the research folder snapshot:
- git status path (resolve) takes 0.8s
- git archive path (package & gzip) takes ~8-11s
=> total time is ~9s (warm) and ~12s (cold) , which matches the python speed
the upload path will depend on the network but it scales with the size of the tarball which is 18.9 MB (18M) (compressed from the 75MB of source)
| View at: [DATABRICKS_URL]/jobs/runs/555 | ||
|
|
||
| === the ai_runtime_task carries the snapshot + provenance paths | ||
| >>> print_requests.py //api/2.2/jobs/runs/submit |
| // ai_runtime_task; a run with no code_source leaves them empty. | ||
| var snap snapshotResult | ||
| if cfg.CodeSource != nil { | ||
| snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir) |
There was a problem hiding this comment.
You also need to check whether cfg.CodeSource.Snapshot exists
…nce fields The air-cli merge migrated the submit path to the typed SDK (jobs.SubmitRun / jobs.AiRuntimeTask + w.Jobs.Submit), but left buildSubmitPayload setting GitStatePath/GitDiffPath and referencing an undropped `snap` — neither compiles, which is why lint, task test, and validate-generated all failed (the package would not build). Fix: restore the `snap` parameter, set only CodeSourcePath, and drop the two git provenance fields from the payload. Update the code_source test to the typed SDK types (jobs.SubmitRun / jobs.SubmitRunResponse) and stop asserting the two dropped fields. TEMP FIX — why this is safe, and the downstream effects: - The typed jobs.AiRuntimeTask has no git_state_path / git_diff_path fields, and neither does its source proto (elastic-spark-common/.../ai_runtime_task.proto). So the typed SDK path cannot carry them. - This is not a functional regression: git_state_path has ZERO consumers anywhere in universe, and the proto never declared the fields — so even the Python CLI's raw-JSON values were silently dropped on backend deserialization. Nothing ever read them. - code_source_path (the field that actually delivers the code to the workload) is in the proto and SDK and still ships, fully typed. - The git_state.json / git_diff.patch sidecar FILES are still uploaded next to the tarball (see snapshot.go) for human inspection; only the payload path pointers are omitted. - Follow-up: if the backend adds git_state_path/git_diff_path to the AiRuntimeTask proto, regenerate the SDK and wire snap.GitStatePath/GitDiffPath back into buildSubmitPayload (the code spot is marked with a TODO). Co-authored-by: Isaac
88259f3 to
c88b599
Compare
The typed-SDK submit migration changed the runs/submit request body: JSON keys are now alphabetized and git_state_path is no longer emitted (the field was dropped — see the buildSubmitPayload TEMP note). Regenerate the run-submit golden output to match, and retitle the section from "snapshot + provenance paths" to "code_source_path" since provenance paths are no longer on the payload. Co-authored-by: Isaac
Changes
Ports code_source snapshotting from the Python air CLI to Go and wires it into air run. When a config sets code_source.type: snapshot, the CLI now packages the local repo into a tarball, uploads it to the workspace (or a Volume), records git provenance sidecars, and attaches the paths (code_source_path, git_state_path, git_diff_path) to the ai_runtime_task payload.
Git-archive tarballs are cached by (commit, include_paths), so re-submitting the same commit skips packaging and upload.
Why
air run needs to ship the user's code to the training workload. This is a local-only port: we archive the local checkout rather than cloning from a remote, since that remote-fetch path is being deprecated. Consequences of that decision, which aren't obvious from the code alone:
Tests
inspecting tar:

